# Libraries Used

import os
import sys

import numpy as np # Numerical computing
import pandas as pd # Data manipulation

import seaborn as sns # R-like high-level plots
import matplotlib.pyplot as plt # Python's base plotting 

# home path
sys.path.append('../')
from src import utils, plots

# Functionality to Process the Data (Weight, Normalize, Model, Correlate, Cluster)
from ProteoForge import normalize, impute
# Functionality to Predict Discordant Peptides (id)
from ProteoForge import correct, weight, model
# Functionality to Cluster Peptides for Grouping (group)
from ProteoForge import corluster, disluster
# Functionality to Generate dPFs (differential proteoforms)
from ProteoForge import proteoform_classifier as pfc
# Initialize the timer
startTime = utils.getTime()

## Configure Plotting Settings
def_colors = [
    "#139593", "#fca311", "#e54f2a",
    "#c3c3c3", "#555555",
    "#690000", "#5f4a00", "#004549"
]

condition_colors = {
    'True Hypoxia (1% O2) - 48 Hr': '#e9c46a',
    'True Hypoxia (1% O2) - 72 Hr': '#f4a261',
    'Normoxia (21% O2) - 48 Hr': '#ade8f4',
    'Normoxia (21% O2) - 72 Hr': '#0077b6',
    # Won't be used since the manuscript didn't cover these conditions
    # 'Chemical Hypoxia (CoCl2) - 48 Hr': '#c77dff',
    # 'Chemical Hypoxia (CoCl2) - 72 Hr': '#8338ec',
    # 'Oxidative Stress (H2O2) - 48 Hr': '#ff006e',
    # 'Oxidative Stress (H2O2) - 72 Hr': '#fb5607',
}

# Set seaborn style
sns.set_theme(
    style="white",
    context="paper",
    palette=def_colors,
    font_scale=1,
    rc={
        "figure.figsize": (6, 4),
        "font.family": "sans-serif",
        "font.sans-serif": ["Arial", "Ubuntu Mono"],
    }
)

# Figure Saving Settings
figure_formats = ["png", "pdf"]
save_to_folder = True
transparent_bg = True
figure_dpi = 300

## Configure dataframe displaying
pd.set_option('display.float_format', lambda x: '%.4f' % x)
pd.set_option('display.max_rows', 50)
pd.set_option('display.max_columns', 50)
pd.set_option('display.width', 500)  # Set a wider display width


## Printing Settings
verbose = True

plots.color_palette( def_colors, save=False )
plots.color_palette( condition_colors, save=False, name='condition', size=2)

home_path = './'
data_name = "hypoxia"
notebook_name = "02-ProteoForge"
data_path = f"{home_path}data/"
input_path = f"{data_path}cleaned/{data_name}/"
output_path = f"{data_path}results/{data_name}/"
figure_path = f"{home_path}figures/{data_name}/{notebook_name}/"

# Create the output folder if it does not exist
if not os.path.exists(output_path):
    os.makedirs(output_path)

# Create figure folder structure, if needed
if save_to_folder:
    for i in figure_formats:
        cur_folder = figure_path + i + "/"
        if not os.path.exists(cur_folder):
            os.makedirs(cur_folder)

uniprot_data = pd.read_feather(f"{input_path}uniprot_data.feather")
print(uniprot_data.shape)
uniprot_data.head()

protein_data = pd.read_feather(f"{input_path}protein_data.feather")
print(protein_data.shape)
protein_data.head()

info_data = pd.read_feather(f"{input_path}info_data.feather")
print(info_data.shape)
info_data.head()

cleaned_data = pd.read_feather(f"{input_path}cleaned_data.feather")
print(cleaned_data.shape)
cleaned_data.head()

imputed_data = pd.read_feather(f"{input_path}imputed_data.feather")
print(imputed_data.shape)
imputed_data.head()

meta = pd.read_csv(f"{input_path}metadata.csv")
meta

# Build a dictionary to match condition_colors keys: 'Condition - Duration Hr'
condition_to_samples = {}
for condition in meta['Condition'].unique():
    for duration in sorted(meta.loc[meta['Condition'] == condition, 'Duration'].unique()):
        key = f"{condition} - {duration} Hr"
        samples = meta.loc[(meta['Condition'] == condition) & (meta['Duration'] == duration), 'Colname'].tolist()
        condition_to_samples[key] = samples

# samples_to_condition mapping
samples_to_condition = {sample: condition for condition, samples in condition_to_samples.items() for sample in samples}
# samples_to_colors mapping
samples_to_colors = {sample: condition_colors[condition] for sample, condition in samples_to_condition.items()}

test_data = info_data.sort_values(
    by=['Protein', 'peptide_start', 'peptide_end']
)[['unique_id', 'Protein', 'Peptide', 'peptide_start', 'peptide_end']]
# peptide_end and start should be int
test_data['peptide_start'] = test_data['peptide_start'].astype(int)
test_data['peptide_end'] = test_data['peptide_end'].astype(int)
# Generate PeptideID from cum
test_data['PeptideID'] = test_data.groupby('Protein').cumcount() + 1

print()
print("📊 MELTING CLEANED AND IMPUTED DATA FOR ANALYSIS")
# Get the imputed and cleaned (with missing values) data into long format
melt_imputed = imputed_data.reset_index().drop(
    columns='Protein'
).melt(
    id_vars='unique_id',
    var_name='Sample',
    value_name='Intensity'
).set_index(['unique_id', 'Sample'])
melt_cleaned = cleaned_data.reset_index().drop(
    columns='Protein'
).melt(
    id_vars='unique_id',
    var_name='Sample',
    value_name='Intensity(Raw)'
).set_index(['unique_id', 'Sample'])

combined_data = pd.concat(
    [melt_imputed, melt_cleaned],
    axis=1
).reset_index()

print()
print("📊 ADDITIONAL INFO ADDED")
combined_data['log10Intensity'] = np.log10(combined_data['Intensity'])
combined_data['Condition'] = combined_data['Sample'].map(samples_to_condition)
combined_data['isReal'] = combined_data['Intensity(Raw)'].notna()

compMiss_table = pd.DataFrame(index=cleaned_data.index)
for cond, col in condition_to_samples.items():
    compMiss_table[cond] = cleaned_data[col].isna().all(axis=1)

compMiss_table = compMiss_table.reset_index().drop(columns='Protein').melt(
    id_vars='unique_id',
    var_name='Condition',
    value_name='isCompleteMiss'
)
combined_data = combined_data.merge(
    compMiss_table,
    on=['unique_id', 'Condition'],
    how='left'
)
combined_data = combined_data[[
    'unique_id', 'Condition', 'Sample', 
    'isCompleteMiss', 'isReal',
    'Intensity(Raw)', 'Intensity', 'log10Intensity', 
]]

print()
print("📊 COMBINING INFO AND QUANT DATA")

# Merge the test_data with combined_data
test_data = test_data.merge(
    combined_data,
    on='unique_id',
    how='left'
)

# conditions to subset for 
conditions = [
    # # Removed the chemical and oxidative versions since they were not used in the manuscript
    # 'Oxidative Stress (H2O2) - 48 Hr',
    # 'Oxidative Stress (H2O2) - 72 Hr',
    # 'Chemical Hypoxia (CoCl2) - 48 Hr',
    # 'Chemical Hypoxia (CoCl2) - 72 Hr',
    'True Hypoxia (1% O2) - 48 Hr',
    'True Hypoxia (1% O2) - 72 Hr',
    'Normoxia (21% O2) - 48 Hr',
    'Normoxia (21% O2) - 72 Hr',
]

# Update the condition_to_samples to only include the conditions of interest
condition_to_samples = {
    cond: condition_to_samples[cond] for cond in conditions if cond in condition_to_samples
}

test_data = test_data[
    test_data['Condition'].isin(condition_to_samples.keys())
].reset_index(drop=True)

print()
print("📏 NORMALIZING INTENSITIES AGAINST CONTROL CONDITION")
test_data = normalize.against_condition(
    long_data=test_data,
    cond_run_dict=condition_to_samples,
    run_col='Sample',
    index_cols=['Protein', 'PeptideID'],
    norm_against='Normoxia (21% O2) - 72 Hr',
    intensity_col='Intensity',
    is_log2=False,
    norm_intensity_col='AdjIntensity'
)
print()
# print("🛡️ CALCULATING ROBUST INTENSITIES FOR CLUSTERING")
# test_data = normalize.build_robust_intensity(
#     test_data,
#     log_col='log10Intensity',
#     adj_col='AdjIntensity',
#     protein_col='Protein',
#     cond_col='Condition',
#     out_col='RobIntensity',
#     verbose=verbose

# )
print(f"Test Data shape after normalization: {test_data.shape}")
test_data.head()

weights_data = weight.generate_weights_data(
    test_data,
    sample_cols=['Sample'],
    log_intensity_col='log10Intensity',
    adj_intensity_col='AdjIntensity',
    control_condition='Normoxia (21% O2) - 72 Hr',
    condition_col='Condition',
    protein_col='Protein',
    peptide_col='Peptide',
    is_real_col='isReal',
    is_comp_miss_col='isCompleteMiss',
    sparse_imputed_val=10**-5,
    dense_imputed_val=0.5,
    verbose=verbose,
)
test_data['Weight'] = (
    (weights_data['W_Impute'] * 0.90) + 
    (weights_data['W_RevTechVar'] * 0.10)
)

print(f"Test Data shape after weight generation: {test_data.shape}")
test_data.head()

n_jobs = 16
model_to_use = 'wls' # Options: 'ols', 'wls', 'rlm', 'robust', 'quantile', glm
correction = {
    'strategy': 'two-step',
    'methods': ['bonferroni', 'fdr_bh']
}
# Initialize the linear model with parameters
cur_model = model.LinearModel(
    data=test_data,
    protein_col="Protein",
    peptide_col="Peptide",
    cond_col="Condition",
    intensity_col="AdjIntensity",
    weight_col="Weight",
)

test_data = cur_model.run_analysis(
    model_type=model_to_use,
    correction_strategy= correction['strategy'],
    correction_methods=correction['methods'],
    n_jobs=n_jobs
)

if correction['strategy'] == 'two-step' and isinstance(correction['methods'], tuple):
    methods_str = " and ".join(correction['methods'])
    print(f"  Ran weighted least squares (WLS) model with two-step correction")
    print(f"  Used {n_jobs} parallel jobs")
    print(f"  Applied {methods_str} corrections")
else:
    method_str = correction['methods'] if isinstance(correction['methods'], str) else ", ".join(correction['methods'])
    print(f"  Ran weighted least squares (WLS) model with {correction['strategy']} correction")
    print(f"  Used {n_jobs} parallel jobs")
    print(f"  Applied {method_str} correction")
print()

pThr = 0.001
palette = {
    'pval': '#4c72b0',
    'prt_pval': '#dd8452',
    'adj_pval': '#55a868'
}
# --- 3. Plotting ---
# Define the columns to plot and their corresponding titles
plot_configs = {
    'pval': 'Raw p-value',
    'prt_pval': 'p-value Within Protein (bonferroni)',
    'adj_pval': 'Adjusted p-value (fdr_bh)'
}

# Calculate log10 p-values for each column, replacing zeros with minimum nonzero value for stability
for col in plot_configs.keys():
    min_nonzero = test_data[col][test_data[col] > 0].min() if (test_data[col] > 0).any() else 1e-300
    test_data[f'log10_{col}'] = np.log10(test_data[col].replace(0, min_nonzero))

fig, axes = plt.subplots(1, 3, figsize=(20, 7), sharex=True, sharey=True)

for ax, (col, title) in zip(axes, plot_configs.items()):
    significant_mask = test_data[col] < pThr
    significant_hits = significant_mask.sum()
    total_hits = test_data[col].count()
    percent_hits = 100 * significant_hits / total_hits if total_hits > 0 else 0
    # Proteins
    total_proteins = test_data['Protein'].nunique()
    significant_proteins = test_data.loc[significant_mask, 'Protein'].nunique()
    percent_proteins = 100 * significant_proteins / total_proteins if total_proteins > 0 else 0
    # Peptides
    total_peptides = test_data['Peptide'].nunique()
    significant_peptides = test_data.loc[significant_mask, 'Peptide'].nunique()
    percent_peptides = 100 * significant_peptides / total_peptides if total_peptides > 0 else 0

    # Drop NaN values for histogram
    log10_vals = test_data[f'log10_{col}'].dropna()
    if len(log10_vals) == 0:
        ax.text(0.5, 0.5, 'No data available', transform=ax.transAxes, ha='center', va='center', fontsize=14)
        continue

    counts, bins = np.histogram(log10_vals, bins=50)
    counts_log10 = np.log10(np.where(counts == 0, 1, counts))

    ax.bar(bins[:-1], counts_log10, width=np.diff(bins), align='edge', color=palette[col], edgecolor='white', linewidth=0.5, alpha=0.8)

    # Add vertical line for threshold
    x_thr = np.log10(pThr)
    ax.axvline(x=x_thr, color='crimson', linestyle='--', linewidth=2)

    # Add vertical annotation next to the line
    ylim = ax.get_ylim()
    ax.text(
        x_thr, ylim[1]*0.95, 
        f'log10={x_thr:.0f}', 
        color='crimson', rotation=90, 
        va='top', ha='right', 
        fontsize=12, fontweight='bold', 
        bbox=dict(boxstyle='round,pad=0.2', fc='white', ec='crimson', alpha=0.7)
    )

    annotation = (
        f'Hits: {significant_hits:,} / {total_hits:,} ({percent_hits:.2f}%)\n'
        f'Proteins: {significant_proteins:,} / {total_proteins:,} ({percent_proteins:.2f}%)\n'
        f'Peptides: {significant_peptides:,} / {total_peptides:,} ({percent_peptides:.2f}%)'
    )
    ax.text(
        0.05, 0.90,
        annotation,
        transform=ax.transAxes,
        fontsize=12, fontweight='bold',
        ha='left', va='top',
        bbox=dict(boxstyle='round,pad=0.3', fc='white', ec='black', alpha=0.75)
    )

    ax.set_title(title + ' (log10)', fontsize=16, fontweight='bold')
    ax.set_xlabel('log10(p-value)', fontsize=12)
    ax.set_ylabel('log10(Frequency)', fontsize=14)
    ax.grid(axis='y', linestyle='--', alpha=0.6)
    ax.tick_params(axis='x', rotation=45)

axes[1].set_ylabel('')
axes[2].set_ylabel('')

fig.suptitle(
    'Comparison of log10 p-value Distributions',
    fontsize=22,
    fontweight='bold',
    y=1.02
    )

sns.despine()
plt.tight_layout(rect=[0, 0, 1, 0.96])
plots.finalize_plot( 
    fig, show=True, save=save_to_folder, 
    filename='pvalue_distributions_Comb',
    filepath=figure_path,
    formats=figure_formats,
    transparent=transparent_bg,
    dpi=figure_dpi
)

test_data['isSignificant'] = test_data['adj_pval'] < pThr
test_data.head()

clusters = disluster.distance_and_cluster(
    data=test_data,
    protein_col='Protein',
    peptide_col='PeptideID',
    cond_col='Condition',
    quant_col='AdjIntensity',
    pdist_metric='euclidean',
    clustering_params={
        'min_clusters': 1,
        'distance_transform': 'corr',
        'clustering_method': 'hybrid_outlier_cut',
        'linkage_method': 'ward',
        'distance_metric': 'euclidean'
    },
    n_jobs=n_jobs,
    verbose=False
)

test_data = test_data.merge(
    clusters[['Protein', 'PeptideID', 'cluster_label']],
    on=['Protein', 'PeptideID'],
    how='left'
).rename(columns={'cluster_label': 'ClusterID'})

# Count proteins and peptides processed
proteins_processed = test_data['Protein'].nunique()
peptides_processed = test_data['Peptide'].nunique()

print(f"  Performed hierarchical clustering with average linkage and correlation distance")
print(f"  Auto-determined optimal number of clusters (1-25 range)")
print(f"  Calculated clusters with {n_jobs} parallel jobs")
print(f"  Processed {proteins_processed} proteins with {peptides_processed} peptides")

summary_data = test_data[[
    'Protein', 'PeptideID', 'Peptide',
    'peptide_start', 'peptide_end',
    'pval', 'adj_pval',
    'isSignificant', 'ClusterID'
]].drop_duplicates().copy()

summary_data = pfc.process_peptide_data(
    summary_data, 
    protein_col='Protein',
    cluster_col='ClusterID',
    significance_col='isSignificant',
    dpf_col='dPF', 
    validate=False,
    verbose=verbose
)
print(f"  Processed {summary_data['dPF'].sum()} differential proteoforms (dPFs)")

summary_data = protein_data.merge(
    summary_data,
    on='Protein',
    how='left'
)

# Add the dPF to the test_data
test_data = test_data.merge(
    summary_data[['Protein', 'PeptideID', 'dPF']],
    on=['Protein', 'PeptideID'],
    how='left'
)

test_data.to_feather(f"{output_path}test_data_Comb.feather")
summary_data.to_feather(f"{output_path}summary_data_Comb.feather")
uniprot_data.to_feather(f"{output_path}uniprot_data_Comb.feather")

print("Notebook Execution Time:", utils.prettyTimer(utils.getTime() - startTime))